import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class ECALayer(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """
    
    def __init__(self, c1, c2, k_size=5):
        super(ECALayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()

        y = self.avg_pool(x)
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y)

        return x * y.expand_as(x)
    
# Example usage
if __name__ == "__main__":
    model = ECALayer(c1=32, c2=32)
    input_tensor = torch.randn(1, 32, 128, 128) 
    output = model(input_tensor)
    print(output.shape)